#pragma once

#include <atomic>
#include <cassert>
#include <condition_variable>
#include <deque>
#include <functional>
#include <iostream> // For basic error output
#include <memory>   // For std::unique_ptr
#include <mutex>
#include <stdexcept> // For std::runtime_error (optional for error handling)
#include <thread>
#include <vector>

namespace sola // New namespace
{

class ThreadPool {
public:
    using Task = std::function<void()>;

    explicit ThreadPool(size_t numThreads = std::thread::hardware_concurrency())
        : running_(false)
        , maxQueueSize_(0) // Default to unbounded queue
    {
        if (numThreads == 0) {
            numThreads = 1; // Ensure at least one thread if default hardware_concurrency is 0
        }
        // start(numThreads); // Optionally start in constructor, or require explicit start() call
    }

    ~ThreadPool() {
        if (running_) {
            stop(); // Ensure threads are joined
        }
    }

    // Non-copyable, non-movable (for simplicity, can be made movable)
    ThreadPool(const ThreadPool&)            = delete;
    ThreadPool& operator=(const ThreadPool&) = delete;
    // ThreadPool(ThreadPool&&) = default; // If std::thread and unique_ptr are used
    // ThreadPool& operator=(ThreadPool&&) = default;

    // Must be called before start().
    void setMaxQueueSize(size_t maxSize) {
        if (running_) {
            // Or throw an exception, or log a warning
            std::cerr << "Warning: Cannot set max queue size after ThreadPool has started." << std::endl;
            return;
        }
        maxQueueSize_ = maxSize;
    }

    void start(size_t numThreads) {
        if (running_)
            return; // Already started
        if (numThreads == 0) {
            // Or throw std::invalid_argument("Number of threads cannot be zero.");
            std::cerr << "Warning: Number of threads for ThreadPool cannot be zero. Defaulting to 1." << std::endl;
            numThreads = 1;
        }

        running_ = true;
        threads_.reserve(numThreads);
        for (size_t i = 0; i < numThreads; ++i) {
            // Using std::make_unique for exception safety with new
            threads_.emplace_back(std::make_unique<std::thread>(&ThreadPool::workerLoop, this));
        }
    }

    void stop() {
        if (!running_)
            return;

        {
            std::unique_lock<std::mutex> lock(queueMutex_);
            running_ = false; // Signal threads to stop
        }                     // Release lock before notifying

        // Notify all waiting threads to wake up and check the running_ flag
        notEmptyCondition_.notify_all();
        notFullCondition_.notify_all(); // Also wake up any threads waiting because queue was full

        for (auto& thread_ptr : threads_) {
            if (thread_ptr && thread_ptr->joinable()) {
                thread_ptr->join();
            }
        }
        threads_.clear(); // Clear the vector of thread pointers
    }

    size_t queueSize() const {
        std::lock_guard<std::mutex> lock(queueMutex_);
        return taskQueue_.size();
    }

    // Submit a task to the thread pool.
    // This function might block if the queue is full (if maxQueueSize_ > 0).
    void run(Task task) {
        if (!running_ && threads_.empty()) { // Handle case where pool is not started or stopped and has no threads
            // Option 1: Execute immediately in the calling thread
            // task();
            // return;
            // Option 2: Or reject the task
            std::cerr << "Warning: ThreadPool is not running or has no threads. Task not queued." << std::endl;
            return;
        }
        if (!task) { // Do not queue null tasks
            return;
        }

        std::unique_lock<std::mutex> lock(queueMutex_);

        // Wait if the queue is full and the pool is running
        notFullCondition_.wait(lock, [this] {
            return !isFull() || !running_; // Stop waiting if not full or not running
        });

        if (!running_) { // Pool was stopped while waiting
            // Optionally, you could throw an exception or just return
            std::cerr << "Warning: ThreadPool stopped while trying to queue task. Task not queued." << std::endl;
            return;
        }

        assert(!isFull()); // Should not be full if wait condition was met while running

        taskQueue_.push_back(std::move(task));
        lock.unlock(); // Unlock before notifying to reduce contention

        notEmptyCondition_.notify_one(); // Notify one worker thread
    }

private:
    bool isFull() const // Caller must hold queueMutex_
    {
        // This assertion is tricky without THREAD_ANNOTATION_ATTRIBUTE__
        // assert(queueMutex_ is locked by current thread);
        return maxQueueSize_ > 0 && taskQueue_.size() >= maxQueueSize_;
    }

    void workerLoop() {
        // Optional: Thread initialization callback could be called here if needed
        // if (threadInitCallback_) threadInitCallback_();

        while (running_) {
            Task task;
            { // Scope for the lock
                std::unique_lock<std::mutex> lock(queueMutex_);

                // Wait until there's a task or the pool is stopped
                notEmptyCondition_.wait(lock, [this] { return !taskQueue_.empty() || !running_; });

                if (!running_ && taskQueue_.empty()) {
                    // Pool stopped and queue is empty, exit thread
                    return;
                }

                if (!taskQueue_.empty()) {
                    task = std::move(taskQueue_.front());
                    taskQueue_.pop_front();

                    // If maxQueueSize_ > 0, notify threads waiting because queue was full
                    if (maxQueueSize_ > 0) {
                        lock.unlock(); // Unlock before notify_one to avoid waking a thread that immediately re-locks
                        notFullCondition_.notify_one();
                    } else {
                        lock.unlock(); // Still good practice to unlock before using task
                    }
                } else {
                    // Spurious wakeup or pool stopping with empty queue, continue loop to check running_
                    lock.unlock();
                    continue;
                }
            } // Lock released

            if (task) {
                try {
                    task();
                } catch (const std::exception& ex) {
                    std::cerr << "Exception caught in ThreadPool worker: " << ex.what() << std::endl;
                    // Decide on error handling: log, continue, abort, etc.
                } catch (...) {
                    std::cerr << "Unknown exception caught in ThreadPool worker." << std::endl;
                }
            }
        } // while(running_)
    }

private:
    mutable std::mutex      queueMutex_;        // Mutex for taskQueue_, running_, and conditions
    std::condition_variable notEmptyCondition_; // Condition for when the queue is not empty
    std::condition_variable notFullCondition_;  // Condition for when the queue is not full

    std::vector<std::unique_ptr<std::thread>> threads_;
    std::deque<Task>                          taskQueue_;

    size_t            maxQueueSize_;
    std::atomic<bool> running_; // Atomic for safe checking in multiple threads, though most modifications are guarded
    // bool running_; // If all accesses are guarded by queueMutex_, std::atomic might be overkill
    // but for a flag read by multiple threads, atomic is safer.
};

} // namespace sola
